import datetime
import os
import shutil

from coin.base.datetime_util import (
    to_datetime,
    to_timestamp_int,
    iterate_date)
from coin.base.hostname import get_hostname
from coin.exchange.base.proto_logger import (
    ProtoLogProvider,
    gen_log_reader_iter,
    to_time_str)
from coin.flow.archive import ensure_dir_exists
from coin.flow.topic_record import TopicRecordWriter
from coin.proto.coin_request_pb2 import StrategyRequestProto
from coin.support.proto_log.logic.log_info import StratInfo


default_strat_stat_root = os.path.expanduser('~/data/strat_stat_proto_log')
default_symbol_stat_root = os.path.expanduser('~/data/symbol_stat_proto_log')


def gen_strategy_stat_dir(*,
                          log_root, trading_date, strat_request, machine):
  td_str = trading_date.strftime('%Y%m%d')
  stat_dir = os.path.join(log_root, td_str, strat_request.strategy_name, machine)
  return stat_dir


def gen_symbol_stat_dir(*,
                        log_root, trading_date, strat_request, machine,
                        account_request, symbol):
  td_str = trading_date.strftime('%Y%m%d')
  acct_str = '%s_%s_%s' % (
      account_request.market_type, account_request.exchange, account_request.owner)
  stat_dir = os.path.join(
      log_root, td_str, strat_request.strategy_name, machine,  acct_str, symbol)
  return stat_dir


def enumerate_strat_stat_log_into_stratinfos(root_dir: str, trading_date: str):
  strat_infos = []
  td_dir = os.path.join(root_dir, trading_date)
  if os.path.exists(td_dir):
    strategy_names = os.listdir(td_dir)
    for strategy_name in strategy_names:
      strat_dir = os.path.join(td_dir, strategy_name)
      machines = os.listdir(strat_dir)
      for machine in machines:
        strat_infos.append(StratInfo(
            strategy_name=strategy_name, machine=machine, trading_date=trading_date))
  return strat_infos


class StrategyStatProvider(ProtoLogProvider):
  def __init__(self, request, log_root=None, machine=None):
    self._request = request
    self._log_root = log_root or default_strat_stat_root
    self._machine = machine or get_hostname()
    self._writer = None
    self._init = True
    super().__init__()

  def _gen_log_path_and_ensure(self, timestamp):
    trading_date = to_datetime(timestamp).date()
    stat_dir = gen_strategy_stat_dir(
        log_root=self._log_root,
        trading_date=trading_date,
        strat_request=self._request,
        machine=self._machine)
    stat_filename = '%s.%s.%s' % (
        self._request.strategy_name, to_time_str(timestamp), timestamp)
    stat_path = os.path.join(stat_dir, stat_filename)
    if self._init:
      if os.path.isdir(stat_dir):
        shutil.rmtree(stat_dir)
      self._init = False
    ensure_dir_exists(stat_dir)
    return stat_path

  def gen_writer(self, timestamp):
    stat_path = self._gen_log_path_and_ensure(timestamp)
    if self._writer is None or self._writer.filepath != stat_path:
      self._writer = TopicRecordWriter(stat_path)
    return self._writer


class StrategyStatReader(object):
  def _enumerate_stat_logs(self, log_dir):
    if not os.path.isdir(log_dir):
      return []
    files = []
    for filename in os.listdir(log_dir):
      filepath = os.path.join(log_dir, filename)
      if os.path.isfile(filepath):
        files.append(filepath)
    return sorted(files)

  def run_from_archive(self, *,
                       on_log_callback,
                       start_time,
                       end_time,
                       root_dir,
                       machine,
                       strategy_name):
    start_ts = to_timestamp_int(start_time)
    end_ts = to_timestamp_int(end_time)
    dt_range = iterate_date(start_time, end_time + datetime.timedelta(days=1))
    log_paths = []
    for dt in dt_range:
      log_dir = gen_strategy_stat_dir(
          log_root=root_dir,
          trading_date=dt,
          strat_request=StrategyRequestProto(strategy_name=strategy_name),
          machine=machine)
      log_paths += self._enumerate_stat_logs(log_dir)
    for ts, log in gen_log_reader_iter(log_paths):
      if ts < start_ts:
        continue
      if ts > end_ts:
        break
      on_log_callback(ts, log)


class SymbolStatProvider(ProtoLogProvider):
  def __init__(self, strat_request, account_request, symbol,
               log_root=None, machine=None):
    self._strat_request = strat_request
    self._log_root = log_root or default_strat_stat_root
    self._machine = machine or get_hostname()
    self._account_request = account_request
    self._symbol = symbol
    self._writer = None
    self._init = True
    super().__init__()

  def _gen_log_path_and_ensure(self, timestamp):
    trading_date = to_datetime(timestamp).date()
    stat_dir = gen_symbol_stat_dir(
        log_root=self._log_root,
        trading_date=trading_date,
        strat_request=self._strat_request,
        machine=self._machine,
        account_request=self._account_request,
        symbol=self._symbol)
    stat_filename = '%s.%s.%s' % (
        self._strat_request.strategy_name, to_time_str(timestamp), timestamp)
    stat_path = os.path.join(stat_dir, stat_filename)
    if self._init:
      if os.path.isdir(stat_dir):
        shutil.rmtree(stat_dir)
      self._init = False
    return stat_path

  def gen_writer(self, timestamp):
    stat_path = self._gen_log_path_and_ensure(timestamp)
    if self._writer is None or self._writer.filepath != stat_path:
      self._writer = TopicRecordWriter(stat_path)
    return self._writer
